// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/small_vector.h>

#include <pollux/common/base/status.h>
#include <pollux/functions/prestosql/types/ip_address_type.h>
#include <pollux/type/simple_function_api.h>
#include <pollux/type/type.h>

namespace kumo::pollux {
    namespace ipaddress {
        constexpr uint8_t kIPV4Bits = 32;
        constexpr uint8_t kIPV6Bits = 128;
        constexpr int kIPPrefixLengthIndex = 16;
        constexpr int kIPPrefixBytes = 17;
        constexpr auto kIpRowIndex = "ip";
        constexpr auto kIpPrefixRowIndex = "prefix";

        inline melon::Expected<int8_t, Status> tryIpPrefixLengthFromIPAddressType(
            const int128_t &intIpAddr) {
            melon::ByteArray16 addrBytes = {0};
            memcpy(&addrBytes, &intIpAddr, sizeof(intIpAddr));
            std::reverse(addrBytes.begin(), addrBytes.end());
            auto tryV6Addr = melon::IPAddressV6::tryFromBinary(addrBytes);
            if (tryV6Addr.hasError()) {
                return melon::makeUnexpected(
                    threadSkipErrorDetails()
                        ? Status::UserError()
                        : Status::UserError("Received invalid ip address"));
            }

            return tryV6Addr.value().isIPv4Mapped()
                       ? ipaddress::kIPV4Bits
                       : ipaddress::kIPV6Bits;
        }

        inline melon::Expected<std::pair<int128_t, int8_t>, Status>
        tryParseIpPrefixString(melon::StringPiece ipprefixString) {
            // Ensure '/' is present
            if (ipprefixString.find('/') == std::string::npos) {
                return melon::makeUnexpected(
                    threadSkipErrorDetails()
                        ? Status::UserError()
                        : Status::UserError(
                            "Cannot cast value to IPPREFIX: {}", ipprefixString));
            }

            auto tryCdirNetwork = melon::IPAddress::tryCreateNetwork(
                ipprefixString, /*defaultCidr*/ -1, /*applyMask*/ false);

            if (tryCdirNetwork.hasError()) {
                return melon::makeUnexpected(
                    threadSkipErrorDetails()
                        ? Status::UserError()
                        : Status::UserError(
                            "Cannot cast value to IPPREFIX: {}", ipprefixString));
            }

            auto [ip, prefix] = tryCdirNetwork.value();
            if (prefix > ((ip.isIPv4Mapped() || ip.isV4())
                              ? ipaddress::kIPV4Bits
                              : ipaddress::kIPV6Bits)) {
                return melon::makeUnexpected(
                    threadSkipErrorDetails()
                        ? Status::UserError()
                        : Status::UserError(
                            "Cannot cast value to IPPREFIX: {}", ipprefixString));
            }

            auto addrBytes = (ip.isIPv4Mapped() || ip.isV4())
                                 ? melon::IPAddress::createIPv4(ip).mask(prefix).createIPv6().toByteArray()
                                 : melon::IPAddress::createIPv6(ip).mask(prefix).toByteArray();

            std::reverse(addrBytes.begin(), addrBytes.end());

            int128_t intAddr = 0;
            static_assert(sizeof(intAddr) == ipaddress::kIPAddressBytes);
            memcpy(&intAddr, &addrBytes, ipaddress::kIPAddressBytes);
            return std::make_pair(intAddr, prefix);
        }
    }; // namespace ipaddress

    class IPPrefixType : public RowType {
        IPPrefixType()
            : RowType(
                {ipaddress::kIpRowIndex, ipaddress::kIpPrefixRowIndex},
                {IPADDRESS(), TINYINT()}) {
        }

    public:
        static const std::shared_ptr<const IPPrefixType> &get() {
            static const std::shared_ptr<const IPPrefixType> instance{
                new IPPrefixType()
            };

            return instance;
        }

        bool equivalent(const Type &other) const override {
            // Pointer comparison works since this type is a singleton.
            return this == &other;
        }

        const char *name() const override {
            return "IPPREFIX";
        }

        std::string toString() const override {
            return name();
        }

        melon::Dynamic serialize() const override {
            melon::Dynamic obj = melon::Dynamic::object;
            obj["name"] = "Type";
            obj["type"] = name();
            return obj;
        }

        const std::vector<TypeParameter> &parameters() const override {
            static const std::vector<TypeParameter> kEmpty = {};
            return kEmpty;
        }
    };

    MELON_ALWAYS_INLINE bool isIPPrefixType(const TypePtr &type) {
        // Pointer comparison works since this type is a singleton.
        return IPPrefixType::get() == type;
    }

    MELON_ALWAYS_INLINE std::shared_ptr<const IPPrefixType> IPPREFIX() {
        return IPPrefixType::get();
    }

    struct IPPrefixT {
        using type = Row<int128_t, int8_t>;
        static constexpr const char *typeName = "ipprefix";
    };

    using IPPrefix = CustomType<IPPrefixT>;
} // namespace kumo::pollux
